# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: donggu

from experimental.prophet.ops.elemwise_math import isnan
from experimental.prophet.ops.null import is_null_op
from experimental.prophet.ops.control_flow import cond
from experimental.prophet.ops.shared import assign, shared
from experimental.prophet.graph.shape import Shape


# a seq [nan, 1, nan, nan, 2, nan] is transformed to [nan, 1, 1, 1, 2, 2]
def last_valid(x, treat_null_valid=False):
  buffer = shared(x.dtype, shape=x.shape)
  if treat_null_valid:
    new_value = cond(isnan(x), buffer, x)
  else:
    new_value = cond(isnan(x) | is_null_op(x)[0], buffer, x)
  buffer = assign(buffer, new_value)
  return buffer


# null, nan is replaced by default_value
def default(x, default_value):
  return cond(isnan(x) | is_null_op(x)[0], default_value, x)
